#include "VTKExport.h"

#include <fstream>
#include <iomanip>
#include <iostream>

#include "Utility/Algebra/Matrix.h"
#include "Utility/Algebra/Point.h"

#include "DataStructure/Input/Element/CONRODData.h"
#include "DataStructure/Input/Element/CRODData.h"
#include "DataStructure/Input/Element/ElementActiveDOFs.h"
#include "DataStructure/Input/Element/ElementTypes.h"
#include "DataStructure/Input/Property/PRODData.h"

#include "DBManager/IService/Input/Element/IElementService.h"
#include "DBManager/IService/Input/Node/INodeService.h"
#include "DBManager/IService/Input/Property/IPRODService.h"
#include "DBManager/IService/Input/Subcase/ISubcaseService.h"
#include "DBManager/IService/Result/IResultDomainsService.h"
#include "DBManager/IService/Result/Nodal/IDISPLACEMENTService.h"
#include "DBManager/IService/Result/Nodal/ITEMPERATUREService.h"

#include "SGFem/Pre/Assembly/Assembly.h"
#include "SGFem/Pre/Assembly/CoordinateProcessor.h"

using SG::Algebra::Matrix;
using SG::Algebra::Point;
using SG::DataStructure::Common::Real;
using namespace SG::DataStructure::FEM;
using namespace SG::DataStructure::Common;
using namespace SG::DBManager;
using SG::FEMSolver::AssemblyEigen;
using SG::FEMSolver::CoordinateProcessor;

/** \enum VTKCellTypes
 *  VTK 非结构化网格的单元类型
 *
 */
// ToDo: 其他单元类型
enum VTKCellTypes
{
    VTK_VERTEX               = 1,
    VTK_LINE                 = 3,
    VTK_TRIANGLE             = 5,
    VTK_QUAD                 = 9,
    VTK_TETRA                = 10,
    VTK_HEXAHEDRON           = 12,
    VTK_WEDGE                = 13,
    VTK_PYRAMID              = 14,
    VTK_QUADRATIC_EDGE       = 21,
    VTK_QUADRATIC_TRIANGLE   = 22,
    VTK_QUADRATIC_QUAD       = 23,
    VTK_QUADRATIC_TETRA      = 24,
    VTK_QUADRATIC_HEXAHEDRON = 25,
    VTK_QUADRATIC_WEDGE      = 26,
    VTK_QUADRATIC_PYRAMID    = 27
    // ToDo: 其他单元类型
};

/// @brief 获取单元类型 `type` 的节点数
/// 标量单元节点尺寸不在此处获得
/// ELETYPE_ScalarSpring, ELETYPE_ScalarSpring2 等单元应该调用具体数据
std::size_t getElementNodeSize (const ElementTypes& type)
{
    switch (type)
    {
    case ElementTypes::ELETYPE_Conm1:
    case ElementTypes::ELETYPE_Conm2: 
    {
        return 1;
        break;
    }

    case ElementTypes::ELETYPE_Link:
    case ElementTypes::ELETYPE_LinkNoRef:
    case ElementTypes::ELETYPE_Bush:
    case ElementTypes::ELETYPE_Tube:
    case ElementTypes::ELETYPE_Bar:
    case ElementTypes::ELETYPE_Beam: 
    {
        return 2;
        break;
    }
    case ElementTypes::ELETYPE_PlaneTri3:
    case ElementTypes::ELETYPE_MemTria3:
    case ElementTypes::ELETYPE_Allman3:
    case ElementTypes::ELETYPE_ThinPlate3:
    case ElementTypes::ELETYPE_ThickPlate3:
    case ElementTypes::ELETYPE_ShellTria3:
    case ElementTypes::ELETYPE_ShellTria3R:
    case ElementTypes::ELETYPE_Composite3: 
    {
        return 3;
        break;
    }
    case ElementTypes::ELETYPE_PlaneQuad4:
    case ElementTypes::ELETYPE_Shear:
    case ElementTypes::ELETYPE_MemQuad4:
    case ElementTypes::ELETYPE_Allman4:
    case ElementTypes::ELETYPE_ThinPlate4:
    case ElementTypes::ELETYPE_ThickPlate4:
    case ElementTypes::ELETYPE_ShellQuad4:
    case ElementTypes::ELETYPE_ShellQUAD4R:
    case ElementTypes::ELETYPE_Composite4:
    case ElementTypes::ELETYPE_Tetra4: 
    {
        return 4;
        break;
    }
    case ElementTypes::ELETYPE_PlaneTri6:
    case ElementTypes::ELETYPE_ShellTria6: {
        return 6;
        break;
    }
    case ElementTypes::ELETYPE_PlaneQuad8:
    case ElementTypes::ELETYPE_ShellQuad8:
    case ElementTypes::ELETYPE_Hexa8SRI: {
        return 8;
        break;
    }
    case ElementTypes::ELETYPE_Hexa20: {
        return 20;
        break;
    }
    case ElementTypes::ELETYPE_Tetra10: {
        return 10;
        break;
    }
    case ElementTypes::ELETYPE_Penta6SRI: {
        return 6;
        break;
    }
    case ElementTypes::ELETYPE_Penta15SR: {
        return 15;
        break;
    }
    case ElementTypes::ELETYPE_Pyram5R: {
        return 5;
        break;
    }
    // ToDo: 其他单元类型
    default:
        break;
    }

    std::cout << "Can't support element type [" << type << "] in SG::Post::vtkOutput.\n";
    assert (0);
    return 0;
}

/// @brief 获取单元类型 `type` 的 VTK 单元类型 VTKCellTypes
/// 标量单元类型不在此处获得
VTKCellTypes getEleVTKCellType (const ElementTypes& type)
{
    switch (type)
    {

    case ElementTypes::ELETYPE_Conm1:
    case ElementTypes::ELETYPE_Conm2: {
        return VTKCellTypes::VTK_VERTEX;
        break;
    }

    case ElementTypes::ELETYPE_Link:
    case ElementTypes::ELETYPE_LinkNoRef:
    case ElementTypes::ELETYPE_Tube:
    case ElementTypes::ELETYPE_Bush:
    case ElementTypes::ELETYPE_Bar:
    case ElementTypes::ELETYPE_Beam: {
        return VTKCellTypes::VTK_LINE;
        break;
    }
    case ElementTypes::ELETYPE_PlaneTri3:
    case ElementTypes::ELETYPE_MemTria3:
    case ElementTypes::ELETYPE_Allman3:
    case ElementTypes::ELETYPE_ThinPlate3:
    case ElementTypes::ELETYPE_ThickPlate3:
    case ElementTypes::ELETYPE_ShellTria3:
    case ElementTypes::ELETYPE_ShellTria3R:
    case ElementTypes::ELETYPE_Composite3: {
        return VTKCellTypes::VTK_TRIANGLE;
        break;
    }
    case ElementTypes::ELETYPE_PlaneQuad4:
    case ElementTypes::ELETYPE_Shear:
    case ElementTypes::ELETYPE_MemQuad4:
    case ElementTypes::ELETYPE_Allman4:
    case ElementTypes::ELETYPE_ThinPlate4:
    case ElementTypes::ELETYPE_ThickPlate4:
    case ElementTypes::ELETYPE_ShellQuad4:
    case ElementTypes::ELETYPE_ShellQUAD4R:
    case ElementTypes::ELETYPE_Composite4: {
        return VTKCellTypes::VTK_QUAD;
        break;
    }
    case ElementTypes::ELETYPE_Hexa8SRI: {
        return VTKCellTypes::VTK_HEXAHEDRON;
        break;
    }
    case ElementTypes::ELETYPE_Tetra4: {
        return VTKCellTypes::VTK_TETRA;
        break;
    }
    case ElementTypes::ELETYPE_PlaneQuad8:
    case ElementTypes::ELETYPE_ShellQuad8: {
        return VTKCellTypes::VTK_QUADRATIC_QUAD;
        break;
    }
    case ElementTypes::ELETYPE_PlaneTri6:
    case ElementTypes::ELETYPE_ShellTria6: {
        return VTKCellTypes::VTK_QUADRATIC_TRIANGLE;
        break;
    }
    case ElementTypes::ELETYPE_Hexa20: {
        return VTKCellTypes::VTK_QUADRATIC_HEXAHEDRON;
        break;
    }
    case ElementTypes::ELETYPE_Tetra10: {
        return VTKCellTypes::VTK_QUADRATIC_TETRA;
        break;
    }
    case ElementTypes::ELETYPE_Penta6SRI: {
        return VTKCellTypes::VTK_WEDGE;
        break;
    }
    case ElementTypes::ELETYPE_Penta15SR: {
        return VTKCellTypes::VTK_QUADRATIC_WEDGE;
        break;
    }
    case ElementTypes::ELETYPE_Pyram5R: {
        return VTKCellTypes::VTK_PYRAMID;
        break;
    }
    // ToDo: 其他单元类型
    default:
        break;
    }

    std::cout << "Can't support element type [" << type << "] in SG::Post::vtkOutput.\n";
    assert (0);
    return VTKCellTypes (0);
}

/// @brief 单元是否需要写出到vtk
/// @param eleType 
/// @return 
bool isInvisibleEleType (const ElementTypes& eleType)
{
    if (ElementTypes::ELETYPE_Weld == eleType || ElementTypes::ELETYPE_CHBDYP == eleType || ElementTypes::ELETYPE_FAST == eleType || ElementTypes::ELETYPE_CHBDYP == eleType || ElementTypes::ELETYPE_CHBDYG == eleType)
    {
        return true;
    }

    return false;
}

/**
 * @brief 遍历特殊标量单元，更新 单元数量eleSize、单元数据块数据个数 nodeSize
 * 
 * 特殊标量单元节点数量可能是一个，也可能是两个, 需要查询具体单元数据处理
 * 
 * @tparam T1                       单元数据结构类型
 * @param    dbServiceFactory       数据库指针
 * @param[in,out]    eleSize        单元数量（VTK文件中 CELLS 数量）
 * @param[in,out]    nodeSize       元素数量（VTK文件中 CELLS 数据块）
 */
template<typename T1>
void updateSpecialEleInfo (DBServiceFactorySP dbServiceFactory, std::size_t& eleSize, std::size_t& nodeSize)
{
    const auto elementService = dbServiceFactory->get<IElementService>();
    std::size_t nEleI = elementService->Size<T1>(); 
    eleSize += nEleI;

    T1 data;
    for (size_t iEle = 0; iEle < nEleI; ++iEle)
    {
        std::size_t nNode{2};
        elementService->FindByIndex(iEle, _OUT data);  
        if(data.m_g[1] == 0)
        {
            nNode = 1;
        }
        nodeSize += (nNode + 1);
    }

}

/**
 * @brief  写入特殊标量单元节点信息，即VTK文件中CELLS词条下方信息
 * 
 * 单元可能只有一个节点，也可能有两个节点
 * @tparam T1                       单元数据结构类型
 * @param    dbServiceFactory       数据库指针
 * @param    outFile                输出流
 */
template<typename T1>
void writeSpecialEleNode (DBServiceFactorySP dbServiceFactory, std::ofstream& outFile)
{
    const auto elementService = dbServiceFactory->get<IElementService>();
    std::size_t nEleI = elementService->Size<T1>(); 
    
    T1 data;
    std::vector<Id> nodeIds;  // 存储单元的节点ID（1或2个）
    for (size_t iEle = 0; iEle < nEleI; ++iEle)
    {
        nodeIds.clear();
        // 获取单元节点id
        elementService->FindByIndex(iEle, _OUT data);
        nodeIds = {data.m_g[0]};
        const auto& G2 = data.m_g[1];
        if (G2 != 0) 
        {
            nodeIds.push_back(G2);
        }

        // 2. 写入节点数 节点索引
        auto& assembly = AssemblyEigen::GetInstance ();
        outFile << nodeIds.size();  // 节点数（1或2）
        for (const auto& nodeId : nodeIds)
        {
            std::size_t nodeIndex;
            assembly.FindNodeIndex(nodeId, nodeIndex);
            outFile << " " << nodeIndex;
        }
        outFile << "\n";
    }
}

/**
 * @brief                           将单元类型信息写出到输出文件 outFile
 * 
 * 当单元只有一个节点时，输出的类型为 `VTK_VERTEX`; 当单元具有两个节点时，写出单元类型为  `VTK_LINE`     
 * 
 * @tparam T1                       单元数据结构类型
 * @param    dbServiceFactory       数据库指针
 * @param    outFile                输出流
 */
template<typename T1>
void writeSpecialEleCellType(DBServiceFactorySP dbServiceFactory, std::ofstream& outFile)
{
    const auto elementService = dbServiceFactory->get<IElementService>();
    std::size_t nEleI = elementService->Size<T1>(); 
    T1 data;
    for (size_t iEle = 0; iEle < nEleI; ++iEle)
    {
        elementService->FindByIndex(iEle, _OUT data);
        VTKCellTypes vtkType = (data.m_g[1] != 0) ? VTK_LINE : VTK_VERTEX;
        outFile << vtkType << "\n";
    }
}

/// @brief 写出模型信息
/// @param dbServiceFactory
/// @param outFile
/// @param discription
void writeModel (const SG::DBManager::DBServiceFactorySP dbServiceFactory, std::ofstream& outFile, const std::string& discription)
{
    const auto nodeService       = dbServiceFactory->get<INodeService> ();
    const auto eleSumInfoService = dbServiceFactory->get<IElementService> ();

    auto nodeSize = nodeService->Size ();
    auto eleTypes = eleSumInfoService->GetAllElementTypes ();

    auto& assembly = AssemblyEigen::GetInstance ();
    if (!assembly.IsInitialized ())
        return;

    auto& coordHandler = CoordinateProcessor::GetInstance ();

    std::size_t eleSize{ 0 };
    std::size_t nEleNode{ 0 };

    for (const auto& eleType : eleTypes)
    {
        if (isInvisibleEleType(eleType))
            continue;

        // 处理标量单元：MASS1、MASS2、ScalarSpring、ScalarSpring2
        if (eleType == ElementTypes::ELETYPE_MASS1)
        {
            updateSpecialEleInfo<CMASS1Data> (dbServiceFactory, eleSize, nEleNode);
            continue;
        }
        if (eleType == ElementTypes::ELETYPE_MASS2)
        {
            updateSpecialEleInfo<CMASS2Data> (dbServiceFactory, eleSize, nEleNode);
            continue;
        }

        if (eleType == ElementTypes::ELETYPE_ScalarSpring)
        {
            updateSpecialEleInfo<CELAS1Data> (dbServiceFactory, eleSize, nEleNode);
            continue;
        }
        if (eleType == ElementTypes::ELETYPE_ScalarSpring2)
        {
            updateSpecialEleInfo<CELAS2Data> (dbServiceFactory, eleSize, nEleNode);
            continue;
        }

        // 非标量单元
        auto tmpEleSize = eleSumInfoService->GetElementSize(eleType);
        eleSize += tmpEleSize;
        nEleNode += tmpEleSize * (getElementNodeSize(eleType) + 1);

    }

    // if (!outFile.open (filename))
    // {
    //     assert (0);
    // }

    // Write header, title, data type
    outFile << "# vtk DataFile Version 3.0\n";
    outFile << discription << "\n";
    outFile << "ASCII\n";

    // Write DATASET
    outFile << "\nDATASET UNSTRUCTURED_GRID\n";
    // POINTS
    outFile << "POINTS " << nodeSize << " float\n";
    NodeData node;
    Point    nodeCoord;
    for (size_t iNode = 0; iNode < nodeSize; ++iNode)
    {
        nodeService->FindByIndex (iNode, _OUT node);
        Id dispCoordId{ 0 };
        coordHandler.GetGlobalCoordinate (node, _OUT nodeCoord, _OUT dispCoordId);

        // bool Deform{ false };
        // if (Deform)
        // {
        //     const auto disp = displacement[iNode];
        //     Real       lambda{ 1.0 };
        //     outFile << std::scientific << std::setw (14) << (nodeCoord.m_x + lambda * disp.m_X) << std::setw (14)
        //             << (nodeCoord.m_y + lambda * disp.m_Y) << std::setw (14) << (nodeCoord.m_z + lambda * disp.m_Z) << "\n";
        // }
        // else
        // {
        outFile << std::scientific << std::setw (14) << nodeCoord.m_x << std::setw (14) << nodeCoord.m_y << std::setw (14) << nodeCoord.m_z
                << "\n";
        // }
    }

    // CELLS
    outFile << "CELLS " << eleSize << " " << nEleNode << "\n";
    for (const auto& eleType : eleTypes)
    {
        if (isInvisibleEleType(eleType))
            continue;

        if (eleType == ElementTypes::ELETYPE_MASS1)
        {
            writeSpecialEleNode<CMASS1Data> (dbServiceFactory, outFile);
            continue;
        }
        if (eleType == ElementTypes::ELETYPE_MASS2)
        {
            writeSpecialEleNode<CMASS2Data> (dbServiceFactory, outFile);
            continue;
        }

        if (eleType == ElementTypes::ELETYPE_ScalarSpring)
        {
            writeSpecialEleNode<CELAS1Data> (dbServiceFactory, outFile);
            continue;
        }
        if (eleType == ElementTypes::ELETYPE_ScalarSpring2)
        {
            writeSpecialEleNode<CELAS2Data> (dbServiceFactory, outFile);
            continue;
        }


        // 非标量单元
        auto tmpEleSize = eleSumInfoService->GetElementSize(eleType);
        auto eleNodes = eleSumInfoService->GetElementNodes(eleType);
        auto eleNodeSize = eleNodes.GetColSize();
        for (size_t iEle = 0; iEle < tmpEleSize; ++iEle)
        {
            outFile << eleNodeSize;
            for (size_t iEleNode = 0; iEleNode < eleNodeSize; ++iEleNode)
            {
                std::size_t nodeIndex;
                assembly.FindNodeIndex(eleNodes(iEle, iEleNode), nodeIndex);
                outFile << " " << nodeIndex;
            }
            outFile << "\n";
        }
        
    }


    // CELL_TYPES
    outFile << "CELL_TYPES " << eleSize << "\n";
    for (const auto& eleType : eleTypes)
    {
        if (isInvisibleEleType(eleType))
            continue;

        // 处理特殊单元：直接调用模板函数，内部获取单元数量
        if (eleType == ElementTypes::ELETYPE_MASS1)
        {
            writeSpecialEleCellType<CMASS1Data>(dbServiceFactory, outFile);
            continue;
        }
        if (eleType == ElementTypes::ELETYPE_MASS2)
        {
            writeSpecialEleCellType<CMASS2Data>(dbServiceFactory, outFile);
            continue;
        }
        if (eleType == ElementTypes::ELETYPE_ScalarSpring)
        {
            writeSpecialEleCellType<CELAS1Data>(dbServiceFactory, outFile);
            continue;
        }
        if (eleType == ElementTypes::ELETYPE_ScalarSpring2)
        {
            writeSpecialEleCellType<CELAS2Data>(dbServiceFactory, outFile);
            continue;
        }

        // 非标量单元
        auto VTKCellType = getEleVTKCellType(eleType);
        auto tmpEleSize = eleSumInfoService->GetElementSize(eleType);
        for (size_t iEle = 0; iEle < tmpEleSize; ++iEle)
        {
            outFile << VTKCellType << "\n";
        }
    }

}

void writeDisplacement (std::ofstream& outFile, std::size_t nodeSize, std::vector<SG::DataStructure::Post::DISPLACEMENTData>& displacement)
{
    // Write  POINT_DATA: displacement
    outFile << "\nPOINT_DATA " << nodeSize << "\n";
    outFile << "VECTORS displacement float\n";

    // 写出所有点的线位移
    for (auto&& i : displacement)
    {
        if (i.m_id >= 101000000)  // CWELD virtual node
        {
            continue;
        }
        outFile << std::scientific << std::setw (14) << (fabs (i.m_X) > 1.0e-8 ? i.m_X : 0.0) << std::setw (14)
                << (fabs (i.m_Y) > 1.0e-8 ? i.m_Y : 0.0) << std::setw (14) << (fabs (i.m_Z) > 1.0e-8 ? i.m_Z : 0.0) << "\n";
    }

    // 写出所有点的角位移
    outFile << "VECTORS displacement_rot float\n";
    for (auto&& i : displacement)
    {
        if (i.m_id >= 101000000)  // CWELD virtual node
        {
            continue;
        }
        outFile << std::scientific << std::setw (14) << (fabs (i.m_RX) > 1.0e-8 ? i.m_RX : 0.0) << std::setw (14)
                << (fabs (i.m_RY) > 1.0e-8 ? i.m_RY : 0.0) << std::setw (14) << (fabs (i.m_RZ) > 1.0e-8 ? i.m_RZ : 0.0) << "\n";
    }
}

void writeTemperature (std::ofstream& outFile, std::size_t nodeSize, std::vector<SG::DataStructure::Post::TEMPERATUREData>& temperatureData)
{
    // Write  POINT_DATA: Temperature
    outFile << "\nPOINT_DATA " << nodeSize << "\n";
    outFile << "SCALARS Temperature float 1\n";
    outFile << "LOOKUP_TABLE default\n";

    // 写出所有点的数据
    for (auto&& i : temperatureData)
    {
        outFile << std::scientific << std::setw (14) << (fabs (i.m_VALUE) > 1.0e-8 ? i.m_VALUE : 0.0) << "\n";
    }
}

void SG::Post::vtkOutput (const SG::DBManager::DBServiceFactorySP dbServiceFactory,
                          const std::string&                      filename,
                          const std::string&                      discription)
{
    const auto nodeService         = dbServiceFactory->get<INodeService> ();
    const auto eleSumInfoService   = dbServiceFactory->get<IElementService> ();
    const auto displacementService = dbServiceFactory->get<IDISPLACEMENTService> ();
    const auto pRodService         = dbServiceFactory->get<IPRODService> ();

    auto nodeSize = nodeService->Size ();
    auto eleTypes = eleSumInfoService->GetAllElementTypes ();

    // 获取所有domain
    std::vector<SG::DataStructure::Common::Id> domainIds = dbServiceFactory->get<IResultDomainsService> ()->GetAllId ();

    for (size_t i = 0; i < domainIds.size (); i++)
    {
        const SG::DataStructure::Common::Id domainId = domainIds[i];
        // 创建 ofstream
        std::ofstream outFile;  //
        if (domainIds.size () == 1)
        {
            outFile.open (filename);
        }
        else
        {
            // 在filename最后一个点号前插入domainId
            std::string path = filename.substr (0, filename.find_last_of ('.')) + "_" + std::to_string (domainId)
                               + filename.substr (filename.find_last_of ('.'));
            outFile.open (path);
        }

        // 写出模型信息
        writeModel (dbServiceFactory, outFile, discription);
        // check analysis type
        SG::DataStructure::Post::ResultDomains domainData;
        dbServiceFactory->get<IResultDomainsService> ()->FindById (domainId, _OUT domainData);
        switch (domainData.m_Analysis)
        {
            case SG::DataStructure::Post::AnalysisType::ANALYSISTYPE_Heat: 
            {
                std::vector<SG::DataStructure::Post::TEMPERATUREData> temperatureData;
                dbServiceFactory->get<ITEMPERATUREService> ()->Find ({ domainId }, _OUT temperatureData);
                writeTemperature (outFile, nodeSize, temperatureData);
                break;
            }

            default: 
            {
                std::vector<SG::DataStructure::Post::DISPLACEMENTData> displacement;
                displacementService->Find ({ domainId }, _OUT displacement);
                writeDisplacement (outFile, nodeSize, displacement);
                break;
            }
        }

        outFile.close ();
    }
}